# coding=utf-8

import sys

import pexpect
import time
from pexpect import *

from app.modules.common.pipelines import *
from app.modules.common.tasks import *
from app.modules.common.tasks import TaskExcecuteError, TaskParamsError
from app.modules.common.utils import *
from app.modules.database.pipelines import BaseTask

log = logging.getLogger(__name__)


class InitImportCassandraPipeLine(Pipeline):
    def __init__(self, name, group=None, timeout=3600, **kwargs):
        Pipeline.__init__(self, name, group=group, timeout=timeout, **kwargs)

        self.add_task(CheckCassandraImportParams())

        self.add_task(StopCassandraServer())
        self.add_task(ClearCassandraDataPath())
        self.add_task(StartupCassandraServer())

        self.add_task(CheckCassandraProcessTask())
        self.add_task(CleanTempPathTask())
        self.add_task(CopyDumpFileTask())
        self.add_task(GunzipDumpFileTask())
        self.add_task(ReBuildSchemaFileTask())
        self.add_task(DropKeyspaceTask())
        self.add_task(ReCreateKeyspaceTask())
        self.add_task(CreateSchemaTask())
        self.add_task(ImportCassandraDataTask())

        self.add_task(CassandraAfterImportTask())
        self.add_task(CleanTempPathTask())

    def _before_execute(self):
        log.info('start execute _before_execute')
        self.get_log(log).do_up()
        log.info('end execute _before_execute')

    def _after_execute(self):
        self.get_log(log).do_down()


class CheckCassandraImportParams(BaseTask):
    def __init__(self):
        Task.__init__(self, 'CheckImportParams')

    def __do_execute__(self, params={'cassandra_path': None,
                                     'username': None,
                                     'password': None,
                                     'temp_path': None,
                                     'dump_path': None,
                                     'keyspace': None
                                     }):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> CheckImportParams:Entered execute()")

        error_message=[]
        param_valid = True
        if not self.param_exists('cassandra_path', params):
            params['cassandra_path'] = '/usr/local/db/cassandra'

        if not self.param_exists('username', params):
            param_valid = False
            error_message.append('params[username] must not empty')
        else:
            log_proxy.info("=======> params[username] is %s", params['username'])

        if not self.param_exists('password', params):
            param_valid = False
            error_message.append('params[password] must not empty')
        else:
            log_proxy.info("=======> params[password] is %s", params['password'])

        if not self.param_exists('temp_path', params):
            param_valid = False
            error_message.append('params[temp_path] must not empty')
        else:
            log_proxy.info("=======> params[temp_path] is %s", params['temp_path'])

        if not self.param_exists('dump_path', params):
            param_valid = False
            error_message.append('params[dump_path] must not empty')
        else:
            log_proxy.info("=======> params[dump_path] is %s", params['dump_path'])

        if not self.param_exists('after_imp_sql_files', params):
            log_proxy.info("=======> params[after_imp_sql_files] is empty")
        else:
            log_proxy.info("=======> params[after_imp_sql_files] is %s", params['after_imp_sql_files'])

        if not self.param_exists('keyspace', params):
            param_valid = False
            error_message.append('params[keyspace] must not empty')
        else:
            log_proxy.info("=======> params[keyspace] is %s", params['keyspace'])

        if param_valid and not os.path.exists(params['dump_path']):
            param_valid = False
            error_message.append('not found dump path[%s].' % params['dump_path'])

        if param_valid and not os.path.exists(params['temp_path']):
            self.create_path(params['temp_path'])

        if not param_valid:
            log_proxy.error("******** %s", error_message)
            raise TaskParamsError(error_message)

        log_proxy.info("=======> CheckImportParams:Exit execute()")


class StopCassandraServer(Task):
    def __init__(self):
        Task.__init__(self, 'StopCassandraServer')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> StopCassandraServer:Entered execute()")
        cassandra_path = params['cassandra_path']
        try:
            sh = '/bin/bash -c "%s/bin/nodetool stopdaemon"' % cassandra_path

            log_proxy.info('=======> start stop cassandra server, shell : %s', sh)
            cassandra = pexpect.spawn('su - cassandra', timeout=30)
            cassandra.logfile = log_proxy
            cassandra.expect('$', timeout=30)
            cassandra.sendline(sh)
            index = cassandra.expect(['Cassandra has shutdown', 'ConnectException: \'Connection refused\''], timeout=60)
            # output, status = run(sh, withexitstatus=1)
            # log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if index == 0 or index == 1:
                log_proxy.info('=======> stop cassandra server success.')
            else:
                raise TaskExcecuteError('stop cassandra server failure')

        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)
        log_proxy.info("=======> StopCassandraServer:Exit execute()")


class ClearCassandraDataPath(Task):
    def __init__(self):
        Task.__init__(self, 'CleanCassandraDataPath')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> CleanCassandraDataPath:Entered execute()")
        cassandra_path = params['cassandra_path']
        try:
            sh = '/bin/bash -c "rm -rf %s/data/* %s/logs/*"' % (cassandra_path, cassandra_path)

            log_proxy.info('=======> start execute shell : %s', sh)
            output, status = run(sh, withexitstatus=1, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status == 0:
                log_proxy.info('=======> execute shell success.')
            else:
                raise TaskExcecuteError('Failed to clear the Cassandra data directory')

        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('Failed to clear the Cassandra data directory, cause: %s.' % e.message)
        log_proxy.info("=======> CleanCassandraDataPath:Exit execute()")


class StartupCassandraServer(Task):
    def __init__(self):
        Task.__init__(self, 'StartupCassandraServer')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> StartupCassandraServer:Entered execute()")
        cassandra_path = params['cassandra_path']
        try:
            log_proxy.info('=======> startup cassandra server')
            cassandra = pexpect.spawn('su - cassandra', timeout=None)
            cassandra.logfile = log_proxy
            cassandra.expect('$', timeout=30)
            cassandra.sendline('%s/bin/cassandra' % cassandra_path)
            index = cassandra.expect('Node localhost/127.0.0.1 state jump to normal', timeout=600)

            if index == 0:
                log_proxy.info('=======> startup cassandra server success.')
            else:
                raise TaskExcecuteError('startup cassandra server failure')

        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('startup cassandra server failure, cause: %s.' % e.message)
        log_proxy.info("=======> StartupCassandraServer:Exit execute()")


class CheckCassandraProcessTask(Task):
    def __init__(self):
        Task.__init__(self, 'CheckMysqlProcess')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> CheckCassandraProcess:Entered execute()")
        try:
            sh = '/bin/bash -c "ps -ef | grep cassandra | grep -v grep | wc -l"'
            log_proxy.info('=======> start check cassandra process, shell : %s', sh)
            output, status = run(sh, withexitstatus=1, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status == 0 and int(output) != 0:
                log_proxy.info('=======> check cassandra process success.')
            else:
                raise TaskExcecuteError('cassandra process does not exist')

        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)
        log_proxy.info("=======> CheckCassandraProcess:Exit execute()")


class CleanTempPathTask(BaseTask):

    def __init__(self):
        Task.__init__(self, 'cleanTempPathTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> cleanTempPathTask:Entered execute()")

        path = '%s/%s' % (params['temp_path'], params['keyspace'])
        self.create_path(path)

        log_proxy.info('********************************************************')
        files_path = '%s/*.cql %s/*.csv.gz %s/*.csv' % (path, path, path)

        self.remove_files(files_path)

        log_proxy.info("=======> cleanTempPathTask:Exit execute()")


class CopyDumpFileTask(Task):

    def __init__(self):
        Task.__init__(self, 'CopyDumpFileTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> CopyDumpFileTask:Entered execute()")
        path = '%s/%s' % (params['temp_path'], params['keyspace'])
        sh = '/bin/bash -c "cp %s/* %s"' % (params['dump_path'], path)

        log_proxy.info('=======> start execute shell : %s', sh)
        try:
            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        if status != 0:
            raise TaskExcecuteError('copy dump file failure.')
        else:
            log_proxy.info('=======> copy dump file success.')

        log_proxy.info("=======> CopyDumpFileTask:Exit execute()")


class GunzipDumpFileTask(Task):

    def __init__(self):
        Task.__init__(self, 'GunzipDumpFileTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> GunzipDumpFileTask:Entered execute()")
        path = '%s/%s' % (params['temp_path'], params['keyspace'])
        sh = '/bin/bash -c "gunzip %s/*.csv.gz"' % path

        log_proxy.info('=======> start execute shell : %s', sh)
        try:
            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status != 0:
                raise TaskExcecuteError('gunzip dump file failure.')
            else:
                log_proxy.info('=======> gunzip dump file success.')
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        log_proxy.info("=======> GunzipDumpFileTask:Exit execute()")


class ReBuildSchemaFileTask(Task):

    def __init__(self):
        Task.__init__(self, 'ReBuildSchemaFileTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> ReBuildSchemaFileTask:Entered execute()")
        temp_path = '%s/%s' % (params['temp_path'], params['keyspace'])
        sh = '/bin/bash -c \'grep -v \"CREATE KEYSPACE\" %s/00_keyspace_schema.cql > %s/00_keyspace_schema.cql.tmp\'' \
             % (temp_path, temp_path)

        log_proxy.info('=======> start execute shell : %s', sh)
        try:
            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status != 0:
                raise TaskExcecuteError('rebuild schema file failure.')
            else:
                log_proxy.info('=======> rebuild schema file success.')
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        log_proxy.info("=======> ReBuildSchemaFileTask:Exit execute()")


class DropKeyspaceTask(Task):

    def __init__(self):
        Task.__init__(self, 'DropKeyspaceTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> Re-DropKeyspaceTask:Entered execute()")
        cassandra_home = params['cassandra_path']
        username = params['username']
        password = params['password']
        _keyspace = params['keyspace']

        sh = '%s/bin/cqlsh -u %s -p%s --no-color -e \"drop keyspace if exists %s;\"' % (cassandra_home, username, password, _keyspace)

        log_proxy.info('=======> start drop keyspace: %s', sh)
        try:
            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status != 0:
                raise TaskExcecuteError('drop keyspace failure.')
            else:
                log_proxy.info('=======> drop keyspace success.')
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        log_proxy.info("=======> DropKeyspaceTask:Exit execute()")


class ReCreateKeyspaceTask(Task):

    def __init__(self):
        Task.__init__(self, 'ReCreateKeyspaceTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> Re-CreateKeyspaceTask:Entered execute()")
        cassandra_home = params['cassandra_path']
        username = params['username']
        password = params['password']
        _keyspace = params['keyspace']

        replication = '{\'class\': \'SimpleStrategy\', \'replication_factor\': 1}'
        sh = '%s/bin/cqlsh -u %s -p%s --no-color -e \"CREATE KEYSPACE %s with replication=%s;\"' \
             % (cassandra_home, username, password, _keyspace, replication)

        log_proxy.info('=======> start create keyspace: %s', sh)
        try:
            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status != 0:
                raise TaskExcecuteError('create keyspace failure.')
            else:
                log_proxy.info('=======> create keyspace success.')
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        log_proxy.info("=======> Re-CreateKeyspaceTask:Exit execute()")


class CreateSchemaTask(Task):

    def __init__(self):
        Task.__init__(self, 'CreateSchemaTask')

    def __replace_schema__(self, params, temp_path, filename):
        log_proxy = self.get_log(log)
        if os.path.isfile(filename):
            log_proxy.info('=====> start replace schame, file[%s]', filename)
            new_filename = temp_path + '/' + os.path.basename(filename) + '.replace'
            keyspace = str(params['keyspace']).strip()
            new_schema = 'CREATE TABLE %s.' % keyspace
            tbs = open(filename)
            lines = tbs.readlines()
            fp = open(new_filename, 'w')
            for line in lines:
                if line.find('CREATE TABLE ups.') >= 0:
                    line = line.replace('CREATE TABLE ups.', new_schema)
                if line.find('ON ups.') >= 0:
                    line = line.replace('ON ups.', ('ON %s.' % keyspace))
                fp.write(line)
            fp.close()
            tbs.close()
            log_proxy.info('=====> end replace schema')
            return new_filename
        else:
            return None

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> CreateSchemaTask:Entered execute()")
        temp_path = '%s/%s' % (params['temp_path'], params['keyspace'])
        cassandra_home = params['cassandra_path']
        username = params['username']
        password = params['password']
        _keyspace = params['keyspace']
        schema_file = temp_path + '/00_keyspace_schema.cql.tmp'

        if _keyspace == 'ups2':
            schema_file = self.__replace_schema__(params, temp_path, schema_file)

        sh = '/bin/bash -c "%s/bin/cqlsh -u %s -p%s --no-color -k %s < %s"' \
             % (cassandra_home, username, password, _keyspace, schema_file)

        log_proxy.info('=======> start create schema: %s', sh)
        try:
            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)

            if status != 0:
                raise TaskExcecuteError('create schema failure.')
            else:
                log_proxy.info('=======> create schema success.')
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        log_proxy.info("=======> CreateSchemaTask:Exit execute()")


class ImportCassandraDataTask(Task):

    def __init__(self):
        Task.__init__(self, 'ImportCassandraDataTask')

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> ImportCassandraDataTask:Entered execute()")
        temp_path = '%s/%s' % (params['temp_path'], params['keyspace'])
        cassandra_home = params['cassandra_path']
        username = params['username']
        password = params['password']
        keyspace = params['keyspace']

        try:
            files = os.listdir(temp_path)

            for file in files:
                if file.endswith('.csv'):
                    table = file.split('.')[0]
                    file_path = temp_path + '/' + file
                    if len(open(file_path,'rU').readlines()) > 1:
                        sh = '%s/bin/cqlsh -u %s -p%s --no-color -k %s -e \"COPY %s FROM \'%s\' with HEADER=true;\"' \
                             % (cassandra_home,
                                username,
                                password,
                                keyspace,
                                table,
                                file_path)
                        log_proxy.info('=======> execute shell : %s', sh)
                        output, status = run(sh, withexitstatus=1, timeout=None, logfile=log_proxy)
                        log_proxy.info('=======> shell status: %s, output: %s', status, output)

                        if status != 0:
                            raise TaskExcecuteError('import table failure.')
                        else:
                            log_proxy.info('=======> import table success.')
        except BaseException, e:
            log_proxy.error('=======> ' + e.message)
            raise TaskExcecuteError('execute shell failure, cause: %s.' % e.message)

        log_proxy.info("=======> ImportCassandraDataTask:Exit execute()")


class CassandraAfterImportTask(Task):

    def __init__(self):
        Task.__init__(self, 'CassandraAfterImportTask')

    def __replace_keyword__(self, params, temp_path, filename):
        log_proxy = self.get_log(log)
        if os.path.isfile(filename):
            log_proxy.info('=====> start replace and copy import after sql file[%s]', filename)
            _keyspace = params['keyspace']
            local_ip = params['local_ip']
            log_proxy.info('=====> local ipaddress: [%s]', local_ip)
            new_filename = temp_path + '/' + os.path.basename(filename)
            tbs = open(filename)
            lines = tbs.readlines()
            fp = open(new_filename, 'w')
            for line in lines:
                if line.find('${dbname}') >= 0:
                    line = line.replace('${dbname}', _keyspace)
                if line.find('$localIp') >= 0:
                    line = line.replace('$localIp', local_ip)
                if line.find('${version}') >= 0:
                    line = line.replace('${version}', str(long(time.time() * 1000)))
                fp.write(line)
            fp.close()
            tbs.close()
            log_proxy.info('=====> end replace and copy import after sql file')
            return new_filename
        else:
            return None

    def __do_execute__(self, params={}):
        log_proxy = self.get_log(log)
        log_proxy.info('********************************************************')
        log_proxy.info("=======> CassandraAfterImportTask:Entered execute()")

        cassandra_home = params['cassandra_path']
        username = params['username']
        password = params['password']
        keyspace = params['keyspace']
        log_proxy.info('=======> start cassandra after import : %s', keyspace)
        try:
            temp_path = '%s/%s/afterimp' % (params['temp_path'], keyspace)
            sh = 'mkdir -p %s' % temp_path
            log_proxy.info('start execute shell : %s', sh)
            output, status = pexpect.run(sh, withexitstatus=1, timeout=30, logfile=log_proxy)
            log_proxy.info('=======> shell status: %s, output: %s', status, output)
            if status != 0:
                raise TaskExcecuteError('create dir[%s] failure, cause: %s.' % (temp_path, output))

            if 'after_imp_sql_files' in params.keys():
                after_imp_sql_files = params['after_imp_sql_files']
                if after_imp_sql_files and len(after_imp_sql_files) > 0:
                    for sql_file in after_imp_sql_files:
                        sql_file = self.__replace_keyword__(params, temp_path, sql_file)
                        if sql_file:
                            sh = '/bin/bash -c "%s/bin/cqlsh -u %s -p%s --no-color -k %s < %s"' \
                                 % (cassandra_home, username, password, keyspace, sql_file)
                            log_proxy.info('start execute shell : %s', sh)

                            output, status = run(sh, withexitstatus=1, timeout=300, logfile=log_proxy)
                            log_proxy.info('=======> shell status: %s, output: %s', status, output)

                            if status == 0:
                                log_proxy.info('=======> cassandra after import success')
                            else:
                                log_proxy.info('=======> cassandra after import failure, cause: %s.', output)
                                raise TaskExcecuteError('cassandra after import failure, cause: %s.' % output)
                        else:
                            log_proxy.info('=======> after the introduction of the import of the script does not exist')
            else:
                log_proxy.info('=======> after the introduction of the import of the script does not exist')

        except Exception, e:
            log_proxy.info('=======> ' + e.message)
            raise TaskExcecuteError('cassandra database after the introduction of the import of the script failure.')
        log_proxy.info("=======> CassandraAfterImportTask:Exit execute()")


